import os

os.environ["RAY_SERVE_USE_PACK_SCHEDULING_STRATEGY"] = "1"

import time

from ray._raylet import NodeID
from ray.serve._private import default_impl
from ray.serve._private.common import DeploymentID, ReplicaID
from ray.serve._private.config import ReplicaConfig
from ray.serve._private.deployment_scheduler import (
    ReplicaSchedulingRequest,
    SpreadDeploymentSchedulingPolicy,
)
from ray.serve._private.test_utils import MockActorClass, MockClusterNodeInfoCache


def dummy():
    pass


def bench(num_replicas: int, num_nodes: int, cpus_per_node: int):
    d_id = DeploymentID(name="deployment1")

    cache = MockClusterNodeInfoCache()
    for i in range(num_nodes):
        node_id = NodeID.from_random().hex()
        cache.add_node(node_id, {"CPU": cpus_per_node})

    scheduler = default_impl.create_deployment_scheduler(
        cache,
        head_node_id_override="fake-head-node-id",
        create_placement_group_fn_override=None,
    )
    scheduler.on_deployment_created(d_id, SpreadDeploymentSchedulingPolicy())
    scheduler.on_deployment_deployed(
        d_id, ReplicaConfig.create(dummy, ray_actor_options={"num_cpus": 1})
    )

    requests = [
        ReplicaSchedulingRequest(
            replica_id=ReplicaID(unique_id=f"r{i}", deployment_id=d_id),
            actor_def=MockActorClass(),
            actor_resources={"CPU": 1},
            actor_options={},
            actor_init_args=(),
            on_scheduled=lambda *a, **kw: None,
        )
        for i in range(num_replicas)
    ]

    start = time.perf_counter()
    scheduler.schedule(upscales={d_id: requests}, downscales={})
    elapsed = time.perf_counter() - start
    return elapsed


if __name__ == "__main__":
    configs = [
        # (replicas, nodes, cpus_per_node)
        (256, 8, 64),
        (512, 16, 64),
        (1024, 32, 64),
        (2048, 64, 64),
        (4096, 64, 128),
        (8192, 128, 128),
        (16384, 256, 128),
    ]

    print(f"{'replicas':>10} {'nodes':>6} {'cpus/node':>10} {'time (s)':>10}")
    print("-" * 42)
    for num_replicas, num_nodes, cpus in configs:
        elapsed = bench(num_replicas, num_nodes, cpus)
        print(f"{num_replicas:>10} {num_nodes:>6} {cpus:>10} {elapsed:>10.3f}")